#include "sprite_renderer.h"
#include "vortex/core/asset.h"
#include <iostream>

namespace vortex
{
    SpriteBatch::~SpriteBatch()
    {
        glDeleteVertexArrays(1, &vao);
        glDeleteBuffers(1, &vbo);
        glDeleteBuffers(1, &ibo);
    }

    void SpriteBatch::Init(uint32_t max_sprites)
    {
        max_vertices = max_sprites * 4;
        max_indices = max_sprites * 6;

        // create vao and vbo
        glGenVertexArrays(1, &vao);
        glBindVertexArray(vao);

        glGenBuffers(1, &vbo);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        glBufferData(GL_ARRAY_BUFFER, max_vertices * sizeof(SpriteVertex), nullptr, GL_DYNAMIC_DRAW);

        // set vertex attributes
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, pos));
        glEnableVertexAttribArray(0);
        glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, color));
        glEnableVertexAttribArray(1);
        glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, uv));
        glEnableVertexAttribArray(2);
        glVertexAttribPointer(3, 1, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, texture_slot));
        glEnableVertexAttribArray(3);

        glGenBuffers(1, &ibo);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, max_indices * sizeof(uint16_t), nullptr, GL_DYNAMIC_DRAW);
    }

    void SpriteBatch::Begin()
    {
        // map buffer
        glBindVertexArray(vao);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        vertex_base = (SpriteVertex*)glMapBuffer(GL_ARRAY_BUFFER, GL_WRITE_ONLY);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        index_base = (uint16_t*)glMapBuffer(GL_ELEMENT_ARRAY_BUFFER, GL_WRITE_ONLY);
        num_vertices = 0;
        num_indices = 0;
        textureid2slot.clear();
    }

    void SpriteBatch::End()
    {
        glUnmapBuffer(GL_ARRAY_BUFFER);
        vertex_base = nullptr;
        glUnmapBuffer(GL_ELEMENT_ARRAY_BUFFER);
        index_base = nullptr;
    }

    bool SpriteBatch::AddSprite(const ColoredQuad& sprite, const glm::mat4& model_transform)
    {
        if (num_indices + sprite.NumIndices() > max_indices)
            return false;
        if (num_vertices + sprite.NumVertices() > max_vertices)
            return false;

        // TODO: add model transformation.

        // copy vertex data
        SpriteVertex* data = vertex_base + num_vertices;
        data[0].pos = { sprite.pos.x, sprite.pos.y };
        data[0].color = sprite.color;
        data[0].uv = glm::vec2(0.0f, 0.0f);
        data[0].texture_slot = -1.0f;

        data[1].pos = { sprite.pos.x, sprite.pos.y + sprite.size.y };
        data[1].color = sprite.color;
        data[1].uv = glm::vec2(0.0f, 0.0f);
        data[1].texture_slot = -1.0f;

        data[2].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y + sprite.size.y };
        data[2].color = sprite.color;
        data[2].uv = glm::vec2(0.0f, 0.0f);
        data[2].texture_slot = -1.0f;

        data[3].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y };
        data[3].color = sprite.color;
        data[3].uv = glm::vec2(0.0f, 0.0f);
        data[3].texture_slot = -1.0f;

        

        // copy index data
        uint16_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;

        num_vertices += 4;
        num_indices += 6;
        return true;
    }

    bool SpriteBatch::AddSprite(const SpriteQuad& sprite, const glm::mat4& model_transform)
    {
        if (num_indices + sprite.NumIndices() >= max_indices)
            return false;
        if (num_vertices + sprite.NumVertices() >= max_vertices)
            return false;
        if (textureid2slot.size() >= max_texture_slots)
            return false;

        GLuint texture_id = sprite.texture->GetId();
        uint32_t texture_slot = 0;
        if (textureid2slot.find(texture_id) == textureid2slot.end())
        {
            // texture id not found, create new
            texture_slot = static_cast<uint32_t>(textureid2slot.size());
            textureid2slot.insert(std::make_pair(texture_id, texture_slot));
        }
        else
        {
            texture_slot = textureid2slot[texture_id];
        }
        // std::cout << "texture slot: " << texture_slot << std::endl;
        // copy vertex data
        SpriteVertex* data = vertex_base + num_vertices;
        data[0].pos = { sprite.pos.x, sprite.pos.y };
        data[0].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[0].uv = glm::vec2(0.0f, 0.0f);
        data[0].texture_slot = (float)texture_slot;

        data[1].pos = { sprite.pos.x, sprite.pos.y + sprite.size.y };
        data[1].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[1].uv = glm::vec2(0.0f, 1.0);
        data[1].texture_slot = (float)texture_slot;

        data[2].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y + sprite.size.y };
        data[2].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[2].uv = glm::vec2(1.0f, 1.0f);
        data[2].texture_slot = (float)texture_slot;

        data[3].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y };
        data[3].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[3].uv = glm::vec2(1.0f, 0.0f);
        data[3].texture_slot = (float)texture_slot;

        // copy index data
        uint16_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;

        num_vertices += 4;
        num_indices += 6;
        return true;
    }

    bool SpriteBatch::AddSprite(const AnimatedSprite& sprite, const glm::mat4& model_transform)
    {
        if (num_indices + sprite.NumIndices() >= max_indices)
            return false;
        if (num_vertices + sprite.NumVertices() >= max_vertices)
            return false;
        if (textureid2slot.size() >= max_texture_slots)
            return false;

        GLuint texture_id = sprite.texture->GetId();
        uint32_t texture_slot = 0;
        if (textureid2slot.find(texture_id) == textureid2slot.end())
        {
            // texture id not found, create new
            texture_slot = static_cast<uint32_t>(textureid2slot.size());
            textureid2slot.insert(std::make_pair(texture_id, texture_slot));
        }
        else
        {
            texture_slot = textureid2slot[texture_id];
        }
        // std::cout << "texture slot: " << texture_slot << std::endl;
        // copy vertex data
        SpriteVertex* data = vertex_base + num_vertices;
        data[0].pos = { sprite.pos.x, sprite.pos.y };
        data[0].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[0].uv = glm::vec2(sprite.uv.x, sprite.uv.y);
        data[0].texture_slot = (float)texture_slot;

        data[1].pos = { sprite.pos.x, sprite.pos.y + sprite.size.y };
        data[1].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[1].uv = glm::vec2(sprite.uv.x, sprite.uv.w);
        data[1].texture_slot = (float)texture_slot;

        data[2].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y + sprite.size.y };
        data[2].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[2].uv = glm::vec2(sprite.uv.z, sprite.uv.w);
        data[2].texture_slot = (float)texture_slot;

        data[3].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y };
        data[3].color = glm::vec4(0.0f, 0.0f, 0.0f, 0.0f);
        data[3].uv = glm::vec2(sprite.uv.z, sprite.uv.y);
        data[3].texture_slot = (float)texture_slot;

        // copy index data
        uint16_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;

        num_vertices += 4;
        num_indices += 6;
        return true;
    }

    SpriteRenderer::SpriteRenderer()
    {
        std::map<vortex::ShaderType, std::string> shader_sources;
        std::string vspath = get_asset_path("shaders/sprite.vs");
        std::string fspath = get_asset_path("shaders/sprite.fs");
        shader_sources.insert(std::make_pair(vortex::ShaderType::Vertex, vspath));
        shader_sources.insert(std::make_pair(vortex::ShaderType::Fragment, fspath));
        program = std::make_shared<ShaderProgram>(shader_sources);

        sprite_batch.Init(MAX_QUADS);
    }

    SpriteRenderer::~SpriteRenderer()
    {
    }

    void SpriteRenderer::Begin(const Environment2D& env)
    {
        program->Bind();
        // set uniforms
        glm::mat4 v = env.camera->getViewMatrix();
        glm::mat4 p = env.camera->getProjectionMatrix();
        glm::mat4 vp = p * v;
        program->SetUniformMat4("vp", vp);
        int slots[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
        program->SetTextureArrayUniform("tex", slots, 16);

        sprite_batch.Begin();
    }

    void SpriteRenderer::Draw(const ColoredQuad& sprite, const glm::mat4& model_transform)
    {
        // copy data into buffer
        bool ret = sprite_batch.AddSprite(sprite, model_transform);
        if (!ret)
        {
            FlushBatch();
            sprite_batch.Begin();
            // assume this will succeed, sprite size less than batch size.
            sprite_batch.AddSprite(sprite, model_transform);
        }
    }

    void SpriteRenderer::Draw(const SpriteQuad& sprite, const glm::mat4& model_transform)
    {
        bool ret = sprite_batch.AddSprite(sprite, model_transform);
        if (!ret)
        {
            FlushBatch();
            sprite_batch.Begin();
            sprite_batch.AddSprite(sprite, model_transform);
        }
    }

    void SpriteRenderer::Draw(const AnimatedSprite& sprite, const glm::mat4& model_transform)
    {
        bool ret = sprite_batch.AddSprite(sprite, model_transform);
        if (!ret)
        {
            FlushBatch();
            sprite_batch.Begin();
            sprite_batch.AddSprite(sprite, model_transform);
        }
    }

    void SpriteRenderer::FlushBatch()
    {
        sprite_batch.End();
        uint32_t num_indices = sprite_batch.NumIndices();
        
        // bind texture slots
        int count = 0;
        for (auto& it : sprite_batch.GetSlotMap())
        {
            glActiveTexture(GL_TEXTURE0 + it.second);
            glBindTexture(GL_TEXTURE_2D, it.first);
            count += 0;
        }
        glDrawElements(GL_TRIANGLES, num_indices, GL_UNSIGNED_SHORT, nullptr);
    }

    void SpriteRenderer::End()
    {
        FlushBatch();
        program->Unbind();
    }
}
